#pragma once

#include "TaskQueue.hpp"
#include "TaskSchedulerStrategy.hpp"

#include <thread>
#include <future>

namespace threadpool {

    class TaskExecutor;

    class ThreadPool {

    public:

        // 声明友元类
        friend class TaskExecutor;

        /**
          * @brief  根据指定的任务调度策略、任务执行器构造一个线程池
          * @param  strategy      任务调度策略 默认为DefaultTaskSchedulerStrategy
          * @param  task_executor 任务执行器 默认为std::shared_ptr<TaskExecutor>
          */
        template<typename StrategyType = DefaultTaskSchedulerStrategy>
        explicit ThreadPool(const StrategyType &strategy = DefaultTaskSchedulerStrategy(),
                            std::shared_ptr<TaskExecutor> task_executor = std::make_shared<TaskExecutor>()) :
                queue_(strategy), stop_(false), work_thread_(0), task_executor_(std::move(task_executor)) {}

        virtual ~ThreadPool() { stop(); }

        /**
          * @brief  增加指定数量的线程
          * @param  num 增加的线程数量
          */
        void add(size_t num);

        /**
          * @brief  删除指定数量的空闲线程
          * @param  num 删除的线程数量
          */
        void remove(size_t num);

        /**
          * @brief  关闭线程池并回收所有线程资源
          * @note   如果有工作线程正在执行任务函数需要一直等待到任务函数执行完成
          */
        void stop();

        /**
          * @brief  获取工作线程数量
          * @retval 工作线程数量
          */
        size_t workThreadCount() const { return work_thread_.load(); }

        /**
          * @brief  获取存活线程数量
          * @retval 存活线程数量
          */
        size_t threadCount() const { return workers_.size(); }

        /**
          * @brief  提交一个任务至线程池
          * @param  task 任务对象
          */
        void submit(const std::shared_ptr<Task> &task);

        /**
          * @brief  提交优先级任务
          * @param  priority 优先级
          * @param  func     任务函数
          * @param  ...args  任务函数参数列表
          * @retval 任务函数返回值的future对象
          */
        template<typename Func, typename... Args>
        std::future<typename std::result_of<Func(Args...)>::type>
        submitPriorityTask(int priority, Func &&func, Args &&...args);

        /**
          * @brief  提交定时任务
          * @param  priority 优先级
          * @param  delay    延迟时间
          * @param  func     任务函数
          * @param  ...args  任务函数参数列表
          * @retval 任务函数返回值的future对象
          */
        template<typename Func, typename... Args>
        std::future<typename std::result_of<Func(Args...)>::type>
        submitTimedTask(int priority, std::chrono::milliseconds delay, Func &&func, Args &&...args);

        /**
          * @brief  提交周期任务
          * @param  priority 优先级
          * @param  period   周期间隔
          * @param  func     任务函数
          * @param  ...args  任务函数参数列表
          */
        template<typename Func, typename... Args>
        void submitPeriodicTask(int priority, std::chrono::milliseconds period, Func &&func, Args &&...args);

    private:

        /**
          * @brief  提交线程自杀任务
          * @param  priority 优先级
          * @param  func     任务函数
          * @param  ...args  任务函数参数列表
          * @retval 任务函数返回值的future对象
          */
        template<typename Func, typename... Args>
        std::future<typename std::result_of<Func(Args...)>::type>
        submitSuicideTask(int priority, Func &&func, Args &&...args);

        /**
          * @brief  线程工作函数
          * @note   虚函数可被重写
          */
        virtual void workThread();

    private:

        // 存活线程
        std::vector<std::thread> workers_;

        // 任务队列
        TaskQueue queue_;

        // 线程池关闭标志
        std::atomic<bool> stop_;

        // 工作线程
        std::atomic<size_t> work_thread_;

        // 存活线程互斥锁
        std::mutex mutex_;

        // 任务执行器
        std::shared_ptr<TaskExecutor> task_executor_;

    };

    template<class Func, class... Args>
    std::future<typename std::result_of<Func(Args...)>::type>
    ThreadPool::submitPriorityTask(int priority, Func &&func, Args &&...args) {
        using return_type = typename std::result_of<Func(Args...)>::type;
        auto task = std::make_shared<std::packaged_task<return_type()> >(
                std::bind(std::forward<Func>(func), std::forward<Args>(args)...));
        queue_.push(std::make_shared<PriorityTask>(
                priority,
                [task]() { (*task)(); })
        );
        return task->get_future();
    }

    template<typename Func, typename... Args>
    std::future<typename std::result_of<Func(Args...)>::type>
    ThreadPool::submitTimedTask(int priority, std::chrono::milliseconds delay, Func &&func, Args &&...args) {
        using return_type = typename std::result_of<Func(Args...)>::type;
        auto task = std::make_shared<std::packaged_task<return_type()> >(
                std::bind(std::forward<Func>(func), std::forward<Args>(args)...));
        queue_.push(std::make_shared<TimedTask>(
                priority,
                delay,
                [task]() { (*task)(); })
        );
        return task->get_future();
    }

    template<typename Func, typename... Args>
    void ThreadPool::submitPeriodicTask(int priority, std::chrono::milliseconds period, Func &&func, Args &&...args) {
        queue_.push(std::make_shared<PeriodicTask>(
                priority,
                period,
                std::bind(std::forward<Func>(func), std::forward<Args>(args)...))
        );
    }

    template<typename Func, typename... Args>
    std::future<typename std::result_of<Func(Args...)>::type>
    ThreadPool::submitSuicideTask(int priority, Func &&func, Args &&...args) {
        using return_type = typename std::result_of<Func(Args...)>::type;
        auto task = std::make_shared<std::packaged_task<return_type()> >(
                std::bind(std::forward<Func>(func), std::forward<Args>(args)...));
        queue_.push(std::make_shared<SuicideTask>(
                priority,
                [task]() { (*task)(); })
        );
        return task->get_future();
    }

    void ThreadPool::add(size_t num) {
        std::lock_guard<std::mutex> lock(mutex_);
        for (size_t i = 0; i < num; ++i) {
            workers_.emplace_back(&ThreadPool::workThread, this);
        }
    }

    void ThreadPool::remove(size_t num) {
        for (size_t i = 0; i < std::min(num, workers_.size()); ++i) {
            submitSuicideTask(0, []() {});
        }
    }

    void ThreadPool::stop() {
        stop_.store(true);
        queue_.stop();
        for (auto &worker: workers_) {
            if (worker.joinable()) {
                worker.join();
            }
        }
        workers_.clear();
    }

    void ThreadPool::submit(const std::shared_ptr<Task> &task) {
        queue_.push(task);
    }

    void ThreadPool::workThread() {
        while (!stop_.load()) {
            std::shared_ptr<Task> task = queue_.pop();
            if (task == nullptr) {
                continue;
            }
            work_thread_.fetch_add(1);
            if (task_executor_->defaultBeforeExecute(task, this) == ThreadStatus::ThreadExit) {
                work_thread_.fetch_sub(1);
                break;
            }
            task->execute();
            if (task_executor_->defaultAfterExecute(task, this) == ThreadStatus::ThreadExit) {
                work_thread_.fetch_sub(1);
                break;
            }
            work_thread_.fetch_sub(1);
        }
    }

}